package com.swak.frame.lock;

import java.util.Collections;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

import org.slf4j.Logger;

import com.google.common.collect.Lists;
import com.swak.frame.cache.CacheProxy;
import com.swak.frame.cache.OutputType;
import com.swak.frame.exception.SwakScriptNotFoundException;
import com.swak.frame.exception.ThrowableWrapper;
import com.swak.frame.logger.SwakLoggerFactory;
import com.swak.frame.timer.CycleTask;
import com.swak.frame.timer.HashedWheelTimer;
import com.swak.frame.timer.Timeout;
import com.swak.frame.timer.TimerTask;
import com.swak.frame.ump.EventInfo;
import com.swak.frame.ump.adapter.AlarmEventAdapter;

/**
 * 分布式锁，增加续租，避免业务执行时间超过锁的时间
 * ClassName: RenewalDistributedLock.java
 * 
 * @author colley.ma
 * @date 2021年12月28日 下午2:32:58
 */
public class RenewalDistributedLock implements DistributedLock {
    protected Logger logger = SwakLoggerFactory.getLogger(getClass()); // 日志记录

    private final Map<String, LockEntry> EXPIRATION_RENEWAL = new ConcurrentHashMap<>();

    private CacheProxy cacheProxy;

    /** 是否需要续租 **/
    private boolean renewalEnabled = false;

    // 脚本Lua SHA
    private String[] scriptLuaShaPool = new String[2];

    private AlarmEventAdapter alarmEventAdapter;

    private HashedWheelTimer wheelTimer;

    @Override
    public boolean tryLock(String lockKey, String requestId, long expireTime,
        TimeUnit timeOutUnit) {
        try {
            boolean isLock = cacheProxy.set(lockKey, requestId, expireTime, timeOutUnit, false);
            if (isLock && renewalEnabled) {
                // 此put是表示有任务在进行中,续期任务可以正常执行
                EXPIRATION_RENEWAL.putIfAbsent(requestId,
                    LockEntry.renewal(lockKey, requestId, LockConst.RENEWAL_SCRIPT));
                renewalExpiration(lockKey, requestId, timeOutUnit.toSeconds(expireTime));
            }
            return isLock;
        } catch (Exception e) {
            logger.error("[swak-lock] - tryLock error", e);
            // 是否续租
            if (renewalEnabled) {
                releaseRenewal(lockKey, requestId);
            }
        }
        // 异常获取锁失败
        return false;
    }

    private void releaseRenewal(String lockKey, String requestId) {
        try {
            LockEntry lockEntry = EXPIRATION_RENEWAL.remove(requestId);
            if (Objects.isNull(lockEntry)) {
                return;
            }
            if (Objects.nonNull(lockEntry.getTimeout())) {
                lockEntry.getTimeout().cancel();
            }
            lockEntry = null;
        } catch (Exception e) {
            alarmEventAdapter
                .event(EventInfo.of("分布式锁续期定时任务取消失败！lockKey:{},requestId:{},message:{}", "lock",
                    new Object[] {lockKey, requestId, e.getMessage()}));
            logger.error(
                "[swak-lock] - releaseDistributedLock remove timer task error and lockKey is {}, requestId is {}",
                lockKey, requestId, e);
        }
    }

    private void renewalExpiration(String lockKey, String requestId, long duration) {
        LockEntry lockEntry = EXPIRATION_RENEWAL.get(requestId);
        if (Objects.isNull(lockEntry)) {
            return;
        }
        long delay = duration / 3;
        TimerTask renewalTask = new CycleTask() {
            @Override
            protected void invoke() throws ThrowableWrapper {
                // task执行的时候检查下 lockEntry 是否unLock时被删掉
                logger.info("[swak-lock] - Renewal runing >>>>>> lockKey:{},ttl:{}", lockKey,
                    cacheProxy.ttl(lockKey));
                LockEntry lockEntryOld = EXPIRATION_RENEWAL.get(requestId);
                if (lockEntryOld == null) {
                    cancel();
                    return;
                }
                Boolean exeResult = executeScript(lockEntry, () -> {
                    String scriptLuaSha = scriptLuaShaPool[lockEntry.getScriptTag().ordinal()];
                    String renewTime = Long.toString(TimeUnit.SECONDS.toMillis(duration));
                    return (Boolean) cacheProxy.evalsha(scriptLuaSha,
                        Collections.singletonList(lockKey),
                        Lists.newArrayList(requestId, renewTime), false, OutputType.BOOLEAN);
                });
                if (!exeResult) {
                    cancel();
                    logger.info("[swak-lock] - 分布式锁续期失败！lockKey:{},requestId:{}", lockKey,
                        requestId);
                }
            }
        }.config(delay, TimeUnit.SECONDS, true);
        Timeout timeout = wheelTimer.newTimeout(renewalTask, delay, TimeUnit.SECONDS);
        lockEntry.setTimeout(timeout);
    }

    @Override
    public boolean tryLock(String lockKey, String requestId) {
        return this.tryLock(lockKey, requestId, 300, TimeUnit.SECONDS);
    }

    @Override
    public boolean tryLock(String lockKey, String requestId, long seconds) {
        return this.tryLock(lockKey, requestId, seconds, TimeUnit.SECONDS);
    }

    /**
     * 释放分布式锁
     */
    @Override
    public boolean unLock(String lockKey, String requestId) {
        Boolean result = null;
        try {
            LockEntry lockEntry = LockEntry.unlock(lockKey, requestId, LockConst.UNLOCK_SCRIPT);
            result = executeScript(lockEntry, () -> {
                String scriptLuaSha = scriptLuaShaPool[lockEntry.getScriptTag().ordinal()];
                return (Boolean) cacheProxy.evalsha(scriptLuaSha,
                    Collections.singletonList(lockKey), Collections.singletonList(requestId), false,
                    OutputType.BOOLEAN);
            });
        } catch (Exception e) {
            logger.error("[swak-lock] - unLock error,lockKey:" + lockKey, e);
        } finally {
            if (renewalEnabled) {
                releaseRenewal(lockKey, requestId);
            }
        }
        return Optional.ofNullable(result).orElse(false);
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        // 加载unlock脚本
        scriptLuaShaPool[LockEntry.ScriptTag.UNLOCK.ordinal()] =
            cacheProxy.scriptLoad(LockConst.UNLOCK_SCRIPT);
        if (renewalEnabled) {
            // 加载续租脚本
            scriptLuaShaPool[LockEntry.ScriptTag.RENEWAL.ordinal()] =
                cacheProxy.scriptLoad(LockConst.RENEWAL_SCRIPT);
        }
    }

    private Boolean executeScript(LockEntry lockEntry, Supplier<Boolean> supplier) {
        int retryTime = LockConst.MAX_RETRY_TIME;
        while ((--retryTime) > 0) {
            try {
                return supplier.get();
            } catch (SwakScriptNotFoundException se) {
                try {
                    String scriptSha = cacheProxy.scriptLoad(lockEntry.getScript());
                    scriptLuaShaPool[lockEntry.getScriptTag().ordinal()] = scriptSha;
                } catch (Exception e) {
                    logger.error("[swak-lock] - load lua script error!", e);
                    alarmEventAdapter.event(EventInfo.of("lua脚本无法上传成功，请联系研发进行处理！message:{}", "lock",
                        new Object[] {e.getMessage()}));
                }
            }
        }
        return Boolean.FALSE;
    }

    public void setCacheProxy(CacheProxy cacheProxy) {
        this.cacheProxy = cacheProxy;
    }

    public void setAlarmEventAdapter(AlarmEventAdapter alarmEventAdapter) {
        this.alarmEventAdapter = alarmEventAdapter;
    }

    public void setWheelTimer(HashedWheelTimer wheelTimer) {
        this.wheelTimer = wheelTimer;
    }

    public void setRenewalEnabled(boolean renewalEnabled) {
        this.renewalEnabled = renewalEnabled;
    }

}
